"""
数据恢复脚本
用于恢复基础数据
"""
import sys
import os
from datetime import datetime
import bcrypt

# 添加项目根目录到PATH
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')))

from backend.database.database import SessionLocal
from backend.models.user import User
from backend.wingknife.models import WingknifeProcess, WingknifeExecutionRecord
from backend.data_clean.models import DataCleanTask

def restore_users():
    """恢复用户数据"""
    db = SessionLocal()
    try:
        # 检查是否已有数据
        if db.query(User).count() > 0:
            print("用户表已有数据，跳过恢复")
            return

        # 创建管理员用户
        hashed_password = bcrypt.hashpw("admin123".encode(), bcrypt.gensalt())
        admin_user = User(
            username="admin",
            email="admin@example.com",
            hashed_password=hashed_password.decode(),
            is_active=True,
            is_superuser=True,
            created_at=datetime.now()
        )
        db.add(admin_user)
        db.commit()
        print("成功创建管理员用户")
    except Exception as e:
        db.rollback()
        print(f"恢复用户数据失败: {str(e)}")
    finally:
        db.close()

def restore_wingknife_data():
    """恢复影刀流程数据"""
    db = SessionLocal()
    try:
        # 检查是否已有数据
        if db.query(WingknifeProcess).count() > 0:
            print("影刀流程表已有数据，跳过恢复")
            return

        # 创建示例流程
        process = WingknifeProcess(
            name="示例流程",
            description="这是一个示例流程",
            process_define_id="sample_process",
            parameters={"param1": "value1", "param2": "value2"},
            output_file_path="/path/to/output",
            created_by="admin",
            created_at=datetime.now(),
            updated_at=datetime.now()
        )
        db.add(process)
        db.commit()

        # 创建示例执行记录
        record = WingknifeExecutionRecord(
            process_id=process.id,
            run_id="sample_run",
            parameters={"param1": "value1"},
            status="completed",
            result={"success": True},
            started_at=datetime.now(),
            completed_at=datetime.now()
        )
        db.add(record)
        db.commit()
        print("成功创建示例流程和执行记录")
    except Exception as e:
        db.rollback()
        print(f"恢复影刀数据失败: {str(e)}")
    finally:
        db.close()

def restore_data_clean_tasks():
    """恢复数据清洗任务"""
    db = SessionLocal()
    try:
        # 检查是否已有数据
        if db.query(DataCleanTask).count() > 0:
            print("数据清洗任务表已有数据，跳过恢复")
            return

        # 创建示例任务
        task = DataCleanTask(
            name="示例清洗任务",
            source_type="file",
            file_path="/path/to/source",
            status="completed",
            created_by="admin",
            created_at=datetime.now(),
            updated_at=datetime.now(),
            completed_at=datetime.now(),
            result_file_path="/path/to/result"
        )
        db.add(task)
        db.commit()
        print("成功创建示例数据清洗任务")
    except Exception as e:
        db.rollback()
        print(f"恢复数据清洗任务失败: {str(e)}")
    finally:
        db.close()

def restore_all():
    """恢复所有数据"""
    print("开始恢复数据...")
    restore_users()
    restore_wingknife_data()
    restore_data_clean_tasks()
    print("数据恢复完成")

if __name__ == "__main__":
    restore_all() 